import json
import os
import uuid
from typing import Literal

from langchain.agents import AgentExecutor, create_json_chat_agent
from langgraph.checkpoint.memory import MemorySaver
import langgraph_api as api

from langchain_core.messages import ToolMessage, SystemMessage, HumanMessage
from langgraph.graph import END, StateGraph, START, add_messages
from langchain_community.chat_models import ChatZhipuAI
from langchain_core.prompt_values import PromptValue
from langgraph.prebuilt.chat_agent_executor import AgentState


tools = [
    api.get_table_data,
    api.load_and_filter_data,
    api.calculate_total_energy,
    api.calculate_total_deck_machinery_energy,
    api.query_device_parameter,
    api.get_device_status_by_time_range,
    api.get_table_name,
    api.calculate_energy_consumption,
    api.calculate_total_energy_consumption,
    api.calculate_total_electricity_generation,
    api.calculate_total_fuel_consumption,
    api.get_dict_data,
    # api.calculate_theoretical_power_generation,
         ]

system = ("1. 不要假设或猜测传入函数的参数值。如果用户的描述不明确，直接输出Error。\n"
          "2. 如果用户没有提及明确的年份, 则默认2024年。\n"
          "3. 如果用户问题中只提及了某日, 那说明时间跨度是当日00:00 到 当日 23:59\n"
          "4. 根据用户的问题选择工具进行处理\n"
          "5. 答案可能不止一个, 请根据工具返回的结果找到所有的答案, 例如: 同一天开机的次数可能不止一次\n"
          "6. 不要臆造数据表名, 不确定的时候请使用工具获取表名\n"
          "7. 进行'时长'计算请不要使用小时制加减, 请全部转换到分钟单位再进行相加或相减计算\n"
          "8. 请认真分析用户的问题和工具的描述, 工具返回的结果可能还需要二次加工, 例如: 用户询问'平均能耗', 可以按照每一天分别调用工具, 直到查出所有天的能耗, 再把结果再进行相加后平均值计算; 用户询问'总能耗'说明是对所有能耗进行求和得到最终答案\n"
          "9. 用户的问题中的名词可能与数据表不太相同, 但是意思相近。"
          "例如: 1. '运行时长'的意思是从设备状态从开机到关机之间所经过的时间; 2. 'DP过程'的意思是DP动作的从ON_DP到OFF_DP之间的之间; 3.'推进系统'的意思是'侧推'设备也是推进系统的一种; 4. '摆动次数'的定义, A架摆出到A架摆回算一次; 5.'A架运行效率的意思是 运行时长/一天24小时;"
          "请根据你的理解做出合理回答."
          )


with open("dict.json", "r", encoding="utf-8") as file:
    context_text = str(json.load(file))


class Zhipu(object):
    def __init__(self, tpt: float = 0):

        self.llm = ChatZhipuAI(
            model="glm-4-plus",
            temperature=tpt,
            zhipuai_api_key="dae00bbfa5604ca49303bdc0ed54b95f.KCE7ukYd34wWi0AI"
        )


class Agent:
    def __init__(self,  llm, tools, system, examples, checkpointer=None):
        self.system = system
        self.examples = ""
        graph = StateGraph(AgentState)
        graph.set_entry_point("llm")
        graph.add_node("llm", self.call_openai)
        graph.add_node("action", self.take_action)
        graph.add_node("reflection", self.reflection)

        graph.add_conditional_edges(
            "llm", self.exists_action, {True: "action", False: END}
        )
        graph.add_conditional_edges(
            "action", self.is_tool_call_failed, {True: "reflection", False: "llm"}
        )
        graph.add_edge("reflection", "llm")
        # graph.add_edge("action", "llm")
        self.graph = graph.compile(checkpointer=checkpointer)
        self.tools = {t.name: t for t in tools}
        self.tools_ai = tools
        self.llm = llm
        # self.model = llm.bind_tools(tools)

    def exists_action(self, state: AgentState):
        result = state["messages"][-1]
        return len(result.tool_calls) > 0

    def is_tool_call_failed(self, state: AgentState):
        last_message = state["messages"][-1]
        if "Error in take_action" in last_message.content:
            return True
        else:
            return False

    def call_openai(self, state: AgentState):
        self.model = self.llm.bind_tools(self.tools_ai)
        messages = [SystemMessage(self.system)] + state["messages"]
        message = self.model.invoke(messages)
        return {"messages": [message]}

    def take_action(self, state: AgentState):
        print(f"take_action called!")
        tool_calls = state["messages"][-1].tool_calls
        results = []
        print(f"take_action called with tool_calls: {tool_calls}")
        try:
            for t in tool_calls:
                print(f"Calling: {t}")
                if not t["name"] in self.tools:  # check for bad tool name from LLM
                    print("\n ....bad tool name....")
                    result = "bad tool name, retry"  # instruct LLM to retry if bad
                else:
                    result = self.tools[t["name"]].invoke(t["args"])
                    # print(f"action {t['name']}, result: {result}")
                results.append(
                    ToolMessage(tool_call_id=t["id"], name=t["name"], content=str(result))
                )
            print("Back to the model!")
            return {"messages": results}
        except Exception as e:
            print(f"Error in take_action: {e}")
            return {"messages": [SystemMessage(f"Error in take_action: {e}")]}

    def reflection(self, state: AgentState):
        last_tool_call = state["messages"][-2].tool_calls[-1]
        tool_name = last_tool_call["name"]
        print(f"Reflecting on tool call: {last_tool_call}")
        # Remove the last tool call from the list of available tools
        self.tools.pop(tool_name, None)
        self.tools_ai = [tool for tool in self.tools_ai if tool.name != tool_name]
        # Instruct the LLM to retry with a different tool

        if not self.tools:
            print("All tools have been removed.")
            return {"messages": [SystemMessage("错误: 所有工具都已经删除, 无法直接工具调用功能。")]}
        return {"messages": [SystemMessage("调用出错的工具已经删除, 请重新尝试。")]}

def get_answer(question):
    zhipu = Zhipu()
    memory = MemorySaver()

    agent = Agent(
        llm=zhipu.llm, tools=tools, system=system, examples="", checkpointer=memory
    )
    config = {"configurable": {"thread_id": str(uuid.uuid4())}}
    # The config is the **second positional argument** to stream() or invoke()!
    events = agent.graph.invoke(
        {"messages": [HumanMessage(question)]}, config
    )
    for event in events:
        if "messages" in event:
            return events["messages"][-1].content
        else:
            return "No messages in event!!!"


if __name__ == '__main__':
    zhipu = Zhipu()




    memory = MemorySaver()

    agent = Agent(
        llm=zhipu.llm, tools=tools, system=system, examples="", checkpointer=memory
    )

    folders = ["database_in_use", "data"]
    if any(not os.path.exists(folder) for folder in folders):
        for folder in folders:
            os.makedirs(folder, exist_ok=True)
        import data_process  # for data process using
    else:
        print("所有文件夹均已存在。不再重新预处理数据。")
        print("需要预处理数据，请删除文件夹后重新运行。")

    # user_input = "2024/8/19 下午A架的第一次开机时间（请以XX:XX输出）？"
    # user_input = "统计2024/8/23上午A架的运行时长（以整数分钟输出）？"
    # user_input = "2024/8/24 下午，A架开机发生在折臂吊车开机之前，是否正确？"
    # user_input = "2024/8/24 有什么关键动作是同时发生的？"
    # user_input = "2024/8/24 上午，甲板机械设备的总能耗是多少（单位化成kWh，保留2位小数）？"
    # user_input = "2024/8/24 16:00 ~ 16:30之间，A架发生了哪些关键动作？"
    # user_input = "2024/8/23/上午A架的运行时长和下午A架开机时长相比，哪个长，长多少（以整数分钟输出）？"
    # user_input = "2024/8/23 19:05什么设备在进行什么动作？"
    # user_input = "2024/8/23 DP过程中，侧推的总能耗是多少（单位化成kWh，保留2位小数）？"
    # user_input = "2024/8/23和2024/8/25小艇入水到小艇落座，折臂吊车的总能耗是多少（单位化成kWh，保留2位小数）？"
    # user_input = "2024/8/23 0:00 ~ 2024/8/25 0:00四个发电机中哪个的能耗最大，能耗为多少（单位化成kWh，保留2位小数）？"
    # user_input = "2024/8/23 DP过程中，侧推的总能耗是多少（单位化成kWh，保留2位小数）？"
    # user_input = "2024/8/23 0:00 ~ 2024/8/25 0:00推进系统能耗占总发电量的比例（以%输出，保留2位小数）？"
    # user_input = "2024/8/23 0:00 ~ 2024/8/25 0:00甲板机械能耗占总发电量的比例（以%输出，保留2位小数）？"
    # user_input = "2024/8/23 0:00 ~ 2024/8/25 0:00发电机组的燃油消耗量为多少（单位化成L，保留2位小数）？"
    # user_input = "以征服者落座为标志，2024/8/23 的深海作业A作业结束的时间（请以XX:XX输出）？"
    # user_input = "控制A架左舷的字段名称是？"
    # user_input = "2024/8/23、 2024/8/24和2024/8/25 A架的平均摆动次数是多少次？"

    user_input = "2024/8/23和2024/8/25 平均作业能耗是多久（单位化成kWh，保留2位小数，下放阶段以ON DP和OFF DP为标志，回收阶段以A架开机和关机为标志）？"
    user_input = "5月20日征服者入水A架外摆的最大角度范围分别是多少度，持续了多久？（请以XX:XX输出）？"
    user_input = "数据中有部分时间A架的角度数据出现了异常，请指出开始时间和结束时间（精确到天）"
    user_input = "一号柴油发电机组滑油压力的范围是多少？"
    user_input = "假设柴油的密度为0.8448kg/L，柴油热值为42.6MJ/kg，请计算2024/8/23 0:00 ~ 2024/8/25 0:00的理论发电量（单位化成kWh，保留2位小数）？"
    user_input = "假设柴油的密度为0.8448kg/L，柴油热值为42.6MJ/kg，请计算2024/8/23 0:00 ~ 2024/8/25 0:00柴油机的发电效率（%，保留2位小数）？"
    user_input = "统计2024/8/24-8/30在9点前开始作业的比例（%，保留2位小数）"
    user_input = "统计2024/8/24-8/30征服者在9点前入水的比例（%，保留2位小数）"

    user_input = "2024/8/20深海作业A作业开始的时间（请以XX:XX输出）？"

    user_input = "在2024年8月27日21时09分，一号柴油发电机组转速是多少？"









    config = {"configurable": {"thread_id": "1"}}
    # The config is the **second positional argument** to stream() or invoke()!
    events = agent.graph.stream(
        {"messages": [HumanMessage(user_input)]}, config, stream_mode="values"
    )

    for event in events:
        if "messages" in event:
            event["messages"][-1].pretty_print()

